#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''按照给定的caffe.proto对prototxt文件进行删除操作

    主要是删除一些caffe-quant不支持的算子熟悉
'''

import os 
import sys
from pathlib import Path 
import pdb
import logging

# 设置包搜索路径
p = Path(__file__).parent 
sys.path.append(str(p))
from py_lib.my_exception import MyException

from py_lib.logger_utils import logger_init 
from py_lib.caffe_proto_fix.args_utils import args_parse
args = args_parse()
logger_init(args.verbosity)
from py_lib.utils.caffe_utils import build_layer_list
from py_lib.utils.caffe_utils import write_layer_list

logger = logging.getLogger(__name__)

def main():
    global caffe_pb2
    if args.proto == 'psp-caffe':
        from py_lib.caffe_proto_fix.proto.python import psp_caffe_pb2 as caffe_pb2
    else:
        from py_lib.caffe_proto_fix.proto.python import caffe_pb2
    layerList = build_layer_list(caffe_pb2, args.file)
    clearCnt = __process_layer_list(layerList, args)
    print (f'Cleared {clearCnt} field')
    write_layer_list(caffe_pb2, args.output, layerList)

def __process_layer_list(layerList, args):
    '''只处理Layer下的第一层Field

    '''
    # 清除掉了的Field
    clearCnt = 0
    clearFields = [f.lower() for f in args.clear_field]
    if len(clearFields) == 0:
        return
    for layer in layerList:
        if args.layer_type is not None:
            if layer.type.lower() != args.layer_type.lower():
                continue
        fieldList = layer.ListFields()
        for fieldTuple in fieldList:
            fDesc = fieldTuple[0]
            logger.debug(f'layer {layer.name} field {fDesc.name}')
            if fDesc.name.lower() in clearFields:
                layer.ClearField(fDesc.name)
                clearCnt += 1
    return clearCnt

if __name__ == '__main__':
    try:
        main()
    except MyException as e:
        print(e)
        exit(ERR_FAIL)
